from datetime import datetime
from typing import Dict

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt


def compute_angle_dft(series: pd.Series) -> float:
    """
    使用 DFT 计算相角，返回角度（单位：度）
    """
    x = series.values
    N = len(x)
    n = np.arange(N)
    base = np.exp(-2j * np.pi * n / N)
    X = np.dot(x, base)  # 基波复数分量
    angle = np.angle(X, deg=True)
    return angle % 360

def compute_magnitude_dft(series: pd.Series) -> float:
    """
    使用 DFT 计算基波幅值
    """
    x = series.values
    N = len(x)
    n = np.arange(N)
    base = np.exp(-2j * np.pi * n / N)
    X = np.dot(x, base)
    mag = 2 * np.abs(X) / N   # 基波幅值
    return float(mag)

def fast_rms(series: pd.Series, window: int) -> pd.Series:
    """
    高效滑动 RMS 计算，避免逐窗口 apply
    """
    x2 = series.values ** 2
    cumsum = np.cumsum(np.insert(x2, 0, 0))  # 前缀和
    # 每个窗口的平方和
    sum_sq = cumsum[window:] - cumsum[:-window]
    rms = np.sqrt(sum_sq / window)
    # 对齐索引
    result = pd.Series(rms, index=series.index[window-1:])
    return result.reindex(series.index)  # 保持原长度

def parse_fault_time(time_str):
    # 将最后的 ':160' 转换为 '.160000'
    if time_str.count(':') == 3:
        time_str = time_str.rsplit(':', 1)
        time_str = f"{time_str[0]}.{time_str[1].ljust(6, '0')}"
    return datetime.strptime(time_str, '%Y-%m-%d %H:%M:%S.%f')


def extract_channel_columns(df: pd.DataFrame) -> dict:
    """
    从 DataFrame 列名中匹配三相电流/电压及零序电流/电压通道。
    支持多关键字匹配，返回匹配到的第一个列名。

    Args:
        df: DataFrame，列名包含电流、电压通道名称

    Returns:
        dict: {
            'IA': '列名',
            'IB': '列名',
            'IC': '列名',
            'UA': '列名',
            'UB': '列名',
            'UC': '列名',
            'I0': '列名或None',
            'U0': '列名或None'
        }
    """
    # 定义关键字组，每组内的任一关键字匹配成功即可
    patterns = {
        'IA': ['IA', 'Ia', 'A相', 'PhaseA_I', '电流_A'],
        'IB': ['IB', 'Ib', 'B相', 'PhaseB_I', '电流_B'],
        'IC': ['IC', 'Ic', 'C相', 'PhaseC_I', '电流_C'],
        'UA': ['UA', 'Ua', 'A相', 'PhaseA_U', '电压_A'],
        'UB': ['UB', 'Ub', 'B相', 'PhaseB_U', '电压_B'],
        'UC': ['UC', 'Uc', 'C相', 'PhaseC_U', '电压_C'],
        'I0': ['I0', '零序电流', '3I0', 'Io', 'ZeroSeq_I', 'IN', '电流_N'],
        'U0': ['U0', '零序电压', '3U0', 'Uo', 'ZeroSeq_U', 'UN', '电压_N']
    }

    result = {key: None for key in patterns}

    for key, kws in patterns.items():
        for col in df.columns:
            col_upper = col.upper()
            if any(kw.upper() in col_upper for kw in kws):
                result[key] = col
                break  # 匹配到第一个就停

    return result


def plot_phasors_polar(angles: dict, magnitudes: dict,
                       prefix: str = "I", title: str = "Phasor Diagram"):
    """
    绘制相量图（极坐标表示，区分相位颜色）
    Args:
        angles: 相角字典 { "IA": 角度, "IB": 角度, ... } (单位: 度)
        magnitudes: 幅值字典 { "IA": 幅值, "IB": 幅值, ... }
        prefix: "I" 表示电流, "U" 表示电压
        title: 图标题
    """
    # 预定义颜色映射
    color_map = {
        "A": "tab:red",
        "B": "tab:green",
        "C": "tab:blue"
    }

    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111, polar=True)
    ax.set_title(title, fontsize=13, pad=20)

    for key, angle in angles.items():
        if not key.startswith(prefix):
            continue
        mag = magnitudes.get(key, 1.0)
        theta = np.deg2rad(angle)
        phase = key[-1].upper()  # 获取A/B/C相
        color = color_map.get(phase, "gray")

        ax.arrow(theta, 0, 0, mag,
                 width=0.015*mag,
                 length_includes_head=True,
                 head_width=0.07*mag,
                 color=color,
                 label=key)

        ax.text(theta, mag * 1.1, key, fontsize=10, ha="center", va="center", color=color)

    ax.set_rmax(max(magnitudes.values()) * 1.3 if magnitudes else 1.0)
    ax.grid(True)
    plt.legend(loc="upper right")
    plt.show()